from mpmath import *
import mpmath


def generate_index(s_plus, s_minus):
        if s_plus == matrix([[1],[0]]) and s_minus == matrix([[1],[0]]):
                index = 0
        elif s_plus == matrix([[0],[1]]) and s_minus == matrix([[0],[1]]):
                index = 1
        elif s_plus == matrix([[1],[0]]) and s_minus == matrix([[0],[1]]):
                index = 2
        elif s_plus == matrix([[0],[1]]) and s_minus == matrix([[1],[0]]):
                index = 3
        else:
                print("Internal Error\n")
                exit(-1)
        return index

def generate_vector(index):
        """
                the first item it returns is the s+ and the second is the s-
        """
        if index == 0:
                return (matrix([[1],[0]]), matrix([[1],[0]]))
        elif index == 1:
                return (matrix([[0],[1]]), matrix([[0],[1]]))
        elif index == 2:
                return (matrix([[1],[0]]), matrix([[0],[1]]))
        elif index == 3:
                return (matrix([[0],[1]]), matrix([[1],[0]]))
        else:
                print("Internal Error\n")
                exit(-1)

def generate_two_indice(index):
        """
                the first item it returns is the s+ and the second is the s-
        """
        if index == 0:
                return (1, 1)
        elif index == 1:
                return (-1, -1)
        elif index == 2:
                return (1, -1)
        elif index == 3:
                return (-1, 1)
        else:
                print("Internal Error\n")
                exit(-1)

